package com.tencent.sr.iris.activity.service.activity;

import cn.hutool.core.util.StrUtil;
import com.google.common.collect.Maps;
import com.tencent.sr.iris.activity.common.config.EsConfig;
import com.tencent.sr.iris.activity.common.enums.IrisErrorCode;
import com.tencent.sr.iris.activity.common.util.AssertUtils;
import com.tencent.sr.iris.activity.common.util.ValidateUtils;
import com.tencent.sr.iris.activity.domain.event.dto.EsQueryDTO;
import com.tencent.sr.iris.activity.interfaces.response.SearchAfterPaginationResponse;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.script.Script;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.aggregations.bucket.MultiBucketsAggregation;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.elasticsearch.search.aggregations.metrics.ParsedCardinality;
import org.elasticsearch.search.aggregations.metrics.ParsedScriptedMetric;
import org.elasticsearch.search.aggregations.metrics.ScriptedMetricAggregationBuilder;
import org.elasticsearch.search.aggregations.metrics.Sum;
import org.elasticsearch.search.aggregations.metrics.SumAggregationBuilder;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.collapse.CollapseBuilder;
import org.elasticsearch.search.sort.SortBuilder;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;

@Slf4j
public abstract class EsDataOperateService {

    @Resource
    protected RestHighLevelClient restHighLevelClient;

    @Resource
    protected EsConfig esConfig;

    protected abstract String[] indices();

    /**
     * 分页搜索
     *
     * @param esQueryDTO
     * @param deSerializerFunc
     * @param <T>
     * @return
     */
    public <T> SearchAfterPaginationResponse<T> searchPage(EsQueryDTO esQueryDTO,
                                                           Function<SearchResponse, List<T>> deSerializerFunc) {

        SearchResponse searchResponse = getSearchResponse(esQueryDTO);

        int pageNum = esQueryDTO.getPageNum();
        int pageSize = esQueryDTO.getPageSize();
        AggregationBuilder aggregation = esQueryDTO.getTotalAggregation();

        SearchAfterPaginationResponse<T> response = new SearchAfterPaginationResponse<>();
        response.setPageNum(pageNum);
        response.setPageSize(pageSize);

        if (aggregation != null) {
            ParsedScriptedMetric metric = searchResponse.getAggregations().get(aggregation.getName());
            response.setTotalCount(Math.toIntExact((Integer) metric.aggregation()));
        } else {
            response.setTotalCount(Math.toIntExact(searchResponse.getHits().getTotalHits().value));
        }

        List<T> list = deSerializerFunc.apply(searchResponse);
        response.setDataList(list);
        return response;
    }

    /**
     * 列表搜索
     *
     * @param esQueryDTO
     * @param deSerializerFunc
     * @param <T>
     * @return
     */
    public <T> List<T> searchList(EsQueryDTO esQueryDTO, Function<SearchResponse, List<T>> deSerializerFunc) {
        SearchResponse searchResponse = getSearchResponse(esQueryDTO);
        return deSerializerFunc.apply(searchResponse);
    }

    /**
     * 获取搜索响应
     *
     * @param esQueryDTO
     * @return
     */
    protected SearchResponse getSearchResponse(EsQueryDTO esQueryDTO) {
        ValidateUtils.validate(esQueryDTO);

        int pageNum = esQueryDTO.getPageNum();
        int pageSize = esQueryDTO.getPageSize();

        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder()
                .query(esQueryDTO.getQuery()).from((pageNum - 1) * pageSize).size(pageSize);
        SortBuilder sortBuilder = esQueryDTO.getSortBuilder();
        if (sortBuilder != null) {
            sourceBuilder.sort(sortBuilder);
        }

        AggregationBuilder aggregation = esQueryDTO.getTotalAggregation();
        if (aggregation != null) {
            sourceBuilder.aggregation(aggregation);
        }

        CollapseBuilder collapseBuilder = esQueryDTO.getCollapseBuilder();
        if (collapseBuilder != null) {
            sourceBuilder.collapse(collapseBuilder);
        }

        String[] indices = indices();
        String indicesStr = String.join(",", indices);
        SearchRequest searchRequest = new SearchRequest(indices);

        String routing = esQueryDTO.getRouting();
        if (StrUtil.isNotBlank(routing)) {
            searchRequest.routing(routing);
        }
        sourceBuilder.trackTotalHits(true);
        searchRequest.source(sourceBuilder);

        SearchResponse searchResponse = null;
        try {
            log.info("搜素索引{}，搜索参数：{}", indicesStr, searchRequest);
            searchResponse = restHighLevelClient.search(searchRequest, RequestOptions.DEFAULT);
        } catch (Exception e) {
            log.error("搜素索引{}错误，搜索参数：{}", indicesStr, searchRequest, e);
        }
        AssertUtils.isNotNull(searchResponse, IrisErrorCode.OPERATE_FAIL);

        log.info("搜索索引{}，搜索参数：{}，返回：{}", indicesStr, searchRequest, searchResponse);
        return searchResponse;
    }

    /**
     * 统计
     *
     * @param query
     * @return
     */
    public long esCount(QueryBuilder query, String routing) {
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder().query(query)
                .from(0).size(0);
        return doCount(routing, sourceBuilder).getHits().getTotalHits().value;
    }

    /**
     * 去重统计
     *
     * @param query
     * @param routing
     * @param distinctField
     * @return
     */
    public long esCountDistinct(QueryBuilder query, String distinctField, String routing) {

        String aggregationName = "distinctCount";
        AggregationBuilder aggregation = AggregationBuilders.cardinality(aggregationName)
                .field(distinctField).precisionThreshold(Integer.MAX_VALUE);
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder().query(query)
                .collapse(new CollapseBuilder(distinctField))
                .aggregation(aggregation).from(0).size(0);

        SearchResponse searchResponse = doCount(routing, sourceBuilder);
        ParsedCardinality cardinality = searchResponse.getAggregations().get(aggregationName);

        return cardinality.getValue();
    }

    /**
     * 准确去重统计
     *
     * @param query
     * @param routing
     * @param distinctField
     * @return
     */
    public long esCountDistinctPrecise(QueryBuilder query, String distinctField, String routing) {
        return esCountDistinctPrecise(query, distinctField, routing, null);
    }

    /**
     * 准确去重统计
     *
     * @param query
     * @param routing
     * @param distinctField
     * @param maxPreciseCount 由于通过脚本精确统计会创建一个HashSet，key太多会给es造成压力，需传入一个上限值
     * @return
     */
    public long esCountDistinctPrecise(QueryBuilder query, String distinctField, String routing, Long maxPreciseCount) {

        AggregationBuilder aggregation = getPreciseCountAggregation(distinctField, maxPreciseCount);
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder().query(query)
                .collapse(new CollapseBuilder(distinctField))
                .aggregation(aggregation).from(0).size(0);

        SearchResponse searchResponse = doCount(routing, sourceBuilder);
        ParsedScriptedMetric metric = searchResponse.getAggregations().get(aggregation.getName());

        return Math.toIntExact((Integer) metric.aggregation());
    }

    /**
     * 统计
     *
     * @param routing
     * @param sourceBuilder
     * @return
     */
    private SearchResponse doCount(String routing, SearchSourceBuilder sourceBuilder) {
        String[] indices = indices();
        String indicesStr = String.join(",", indices);
        SearchRequest searchRequest = new SearchRequest(indices);
        searchRequest.source(sourceBuilder).routing(routing);
        sourceBuilder.trackTotalHits(true);
        SearchResponse response = null;
        try {
            log.info("es统计索引{}，参数: {}", indicesStr, searchRequest);
            response = restHighLevelClient.search(searchRequest, RequestOptions.DEFAULT);
        } catch (IOException e) {
            log.error("es统计索引{}错误，参数: {}", indicesStr, searchRequest, e);
        }
        AssertUtils.isNotNull(response, IrisErrorCode.OPERATE_FAIL);

        log.info("es统计索引{}成功，参数: {}，返回：{}", indicesStr, searchRequest, response);
        return response;
    }

    /**
     * 分组统计
     *
     * @param routing
     * @param query
     * @param fieldName
     * @return
     */
    public Map<String, Long> countGroupBy(String routing, QueryBuilder query, String fieldName) {

        Integer termsSize = esConfig.getTermsSize();
        String aggregationName = "groupByAggregation";
        TermsAggregationBuilder aggregation = AggregationBuilders.terms(aggregationName)
                .field(fieldName).size(termsSize).shardSize(termsSize);
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder().query(query)
                .from(0).size(0).aggregation(aggregation);

        SearchResponse searchResponse = doCount(routing, sourceBuilder);
        Terms terms = searchResponse.getAggregations().get(aggregationName);

        return terms.getBuckets().stream().collect(Collectors.toMap(
                MultiBucketsAggregation.Bucket::getKeyAsString,
                i -> i.getDocCount(), (i1, i2) -> i1));
    }

    /**
     * 分组求和，先fieldName分组再求和sumAggField字段
     *
     * @param routing     routing
     * @param query       query
     * @param fieldName   分组字段
     * @param sumAggField 求和字段
     * @return 求和结果
     */
    public Map<String, Long> sumGroupBy(String routing, QueryBuilder query, String fieldName, String sumAggField) {

        Integer termsSize = esConfig.getTermsSize();
        String aggregationName = "groupByAggregation";
        TermsAggregationBuilder aggregation = AggregationBuilders.terms(aggregationName)
                .field(fieldName).size(termsSize).shardSize(termsSize);

        String sumAggName = sumAggField + "_sum";
        SumAggregationBuilder taskSum = AggregationBuilders.sum(sumAggName)
                .field(sumAggField);
        aggregation.subAggregation(taskSum);

        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder().query(query)
                .from(0).size(0).aggregation(aggregation);

        SearchResponse searchResponse = doCount(routing, sourceBuilder);
        Terms terms = searchResponse.getAggregations().get(aggregationName);

        List<? extends Terms.Bucket> buckets = terms.getBuckets();
        Map<String, Long> taskMap = Maps.newHashMapWithExpectedSize(buckets.size());
        for (Terms.Bucket bucket : buckets) {
            // 获取求和
            Aggregations aggregations1 = bucket.getAggregations();
            Sum sum = aggregations1.get(taskSum.getName());
            taskMap.put(bucket.getKeyAsString(), (long) sum.getValue());
        }
        return taskMap;
    }

    /**
     * 通过脚本实现获取精确的去重后的总数，如果总数太大可能会影响es性能
     *
     * @param countKey
     * @param max      由于通过脚本精确统计会创建一个HashSet，key太多会给es造成压力，需传入一个上限值
     * @return
     */
    public AggregationBuilder getPreciseCountAggregation(String countKey, Long max) {
        ScriptedMetricAggregationBuilder totalCount = AggregationBuilders
                .scriptedMetric("totalCount")
                .initScript(new Script("state.nums=new HashSet();"))
                .combineScript(new Script("return state.nums.size();"))
                .reduceScript(new Script(
                        new StringBuilder("long result=0;")
                                .append("for(e in states){")
                                .append("  if(!Objects.isNull(e)){")
                                .append("     result+=e;")
                                .append("  }")
                                .append("}")
                                .append("return result;").toString()
                ));
        if (max == null) {
            totalCount.mapScript(new Script(
                    new StringBuilder("String cKey = doc.").append(countKey).append(".value;")
                            .append("state.nums.add(cKey);").toString()));
        } else {
            String script = new StringBuilder("if(state.nums.size() < ").append(max).append(")")
                    .append("{String cKey = doc.").append(countKey).append(".value;")
                    .append("state.nums.add(cKey);}").toString();
            totalCount.mapScript(new Script(script));
        }
        return totalCount;
    }

}
